/* Author: Sabuj Pattanayek
 * Class: CSCI4330/5330, Dr. R. Butler
 * Date: 
 * Purpose: Runs the n-queens algorithm using MPI. Options are:
 * -q   : Just print the number of solutions and not the solutions themselves.
 * -n N : The dimensions of the board (NxN) and the number of queens to place on the board (N > 0). The current board size limit is given in maxBoardSize.
 *  Note : If you're running this on a single machine uniprocessor, you will get the best performance if you don't run the program through mpiexec and mpd. However, for testing purposes on single machine uniprocessor use X == Y where X and Y are numbers defined in the command "mpiexec -n X ./p2 -q -n Y".
*/

#include <stdio.h>
#include "mpi.h"
#include <unistd.h>
#include <stdlib.h>
#include <string.h>

int totalRanks = 1;
int printSolutions = 1;
int boardSize = 4;
unsigned int numSolutions = 0;
int myRank;

#define maxBoardSize 25
/* 2 characters + 1 space + null character. I don't expect NxN, where N > 99. */
#define maxIntLenInChar 4

/* For printing success or failure at the end */
/* http://www.jsomers.com/nqueen_demo/nqueens.html */
int solutions[] = {1, 0, 0, 2, 10, 4, 40, 92, 352, 724, 2680, 14200, 73712, 365596, 2279184, 14772512, 95815104, 666090624};

void usage(void) {
	if (myRank == 0) {
		printf("Usage: p2 [opts]\n");
		printf("Options: \n\n");
		printf("-q   : Just print the number of solutions and not the solutions themselves.\n");
		printf("-n N : The dimensions of the board (NxN) and the number of queens to place on the board (N > 0). The current limit is a %dx%d board.\n", maxBoardSize, maxBoardSize);
	}
}

void solve(int myWorkRow) {
	int r, c, queen, solutionRows[maxBoardSize], d, hit, numQueens = 1;
	int startCol, startRow, prevRow, goToNextCol;
	char *outSols, oneSol[maxIntLenInChar];
	int solutionSize = boardSize * maxIntLenInChar;
	MPI_Request unusedReq;

	/* The first column of the solution must be the assigned work row */
	solutionRows[0] = myWorkRow;	

	if (boardSize == 1) {
		/* we don't need to lock the mutex since there is only one work row */
		numSolutions++;
		if (printSolutions) {
			outSols = malloc(solutionSize);
			outSols[0] = '\0';
			for (d = 0; d < boardSize; d++) {
				snprintf(oneSol, maxIntLenInChar, "%d ", solutionRows[d]);
				strcat(outSols, oneSol);				
			}
			if (totalRanks > 1) {
				/*printf("Sent solution\n");*/
				MPI_Isend(outSols, solutionSize, MPI_CHAR, 0, 4, MPI_COMM_WORLD, &unusedReq);
			}
			else {
				printf("%s\n", outSols);
			}
		}
		return;
	}

	/* The first column cannot contain another queen, so start from the 2nd column */
	startCol = 1;
	startRow = 0;
	while (startCol != 0) {
		for (c = startCol; c < boardSize; c++) {
			goToNextCol = 0;
			for (r = startRow; r < boardSize; r++) {
				solutionRows[c] = r;
				hit = 0;
				/* Check all rows first */
				for (queen = 0; queen < numQueens; queen++) {
					if (solutionRows[queen] == r) {
						hit = 1;
						break;
					}
				}
				if (hit) {
					continue;
				}
				/* Check the diagonals of the queen in the closest column first. */
				for (queen = (c-1); queen >= 0; queen--) {
					/* Distance to column that is being checked from one of the other columns */
					d = c - queen;
					/* Check forward diagonals */
					/* r + d */
					if (solutionRows[queen] + d == r) {
						hit = 1;
						break;
					}
					/* r - d */
					if (solutionRows[queen] - d == r) {
						hit = 1;
						break;
					}
				}
				if (!hit) {
					/* Did we find a complete solution? */
					if (c + 1 == boardSize) {
						numSolutions++;
						if (printSolutions) {
							outSols = malloc(solutionSize);
							outSols[0] = '\0';
							for (d = 0; d < boardSize; d++) {
								snprintf(oneSol, maxIntLenInChar, "%d ", solutionRows[d]);
								strcat(outSols, oneSol);
							}
							if (totalRanks > 1) {
								/*printf("Sent solution\n");*/
								/*MPI_Isend(outSols, solutionSize, MPI_CHAR, 0, 4, MPI_COMM_WORLD, &unusedReq);*/
								MPI_Send(outSols, solutionSize, MPI_CHAR, 0, 4, MPI_COMM_WORLD);
							}
							else {
								printf("%s\n", outSols);
							}
						}
					}
					else {
						numQueens++;
						goToNextCol = 1;
						break;
					}
				}
			}
			/* Start checking at the next column */
			if (goToNextCol) {
				startCol++;
				startRow = 0;
				break;
			}
			/* If we checked the last row without prematurely breaking out of the loops, then no solution was found in the column. */
			if (r == boardSize) {
				break;
			}
		}
		/* Go back one column and check the remaining rows */
		if ((c == boardSize) || (r == boardSize)) {
			startCol--;
			numQueens--;
			prevRow =  solutionRows[startCol];
			prevRow++;
			startRow = prevRow;
		}
	}
}

void setup() {
	int myWorkRow, rankNum, rowCount, ready = 1;
	int allDone = 0;
	int newSol = 1;
	int solutionSize = boardSize * maxIntLenInChar;
	int receivedWorkMsg, receivedTerminateMsg, receivedReadyMsg, receivedSolMsg;
	unsigned int returnedNumSolutions, receivedSolMsgs = 0;
	unsigned int sentSolMsgs = 0;
	MPI_Request unusedReq, receivedWorkReq, receivedTerminateReq, receivedReadyReq, receivedSol;
	MPI_Status unusedStat, readyStat;
	char *outSols;

	if (totalRanks > 1) {
		/* Multiple processes */
		if (myRank == 0) {
			/* Hand work out to other ranks. */
			/* This bloat is purely to reduce if checks */
			if (printSolutions) {
				outSols = malloc(solutionSize);
				outSols[0] = '\0';
				for (rowCount = 0; rowCount < boardSize; rowCount++) {
					MPI_Irecv(&ready, 1, MPI_INT, MPI_ANY_SOURCE, 1, MPI_COMM_WORLD, &receivedReadyReq);
					/* Loop forever until I receive any ready request */
					for(;;) {
						/*printf ("stuck\n");*/
						MPI_Test(&receivedReadyReq, &receivedReadyMsg, &readyStat);
						if (receivedReadyMsg) {
							/*printf("sent work\n");*/
							MPI_Isend(&rowCount, 1, MPI_INT, readyStat.MPI_SOURCE, 0, MPI_COMM_WORLD, &unusedReq);
							break;
						}
						if (newSol) {
							/*printf("Asked for solution\n");*/
							MPI_Irecv(outSols, solutionSize, MPI_CHAR, MPI_ANY_SOURCE, 4, MPI_COMM_WORLD, &receivedSol);
							sentSolMsgs++;
							newSol = 0;
						}
						MPI_Test(&receivedSol, &receivedSolMsg, &unusedStat);
						if (receivedSolMsg) {
							printf("%s\n", outSols);
							receivedSolMsgs++;
							newSol = 1;
						}
					}
				}

				/* Receive the number of solutions for each row */
				for (rowCount = 0; rowCount < boardSize; rowCount++) {
					MPI_Recv(&returnedNumSolutions, 1, MPI_UNSIGNED, MPI_ANY_SOURCE, 3, MPI_COMM_WORLD, &unusedStat);
					numSolutions += returnedNumSolutions;
					if ((sentSolMsgs > receivedSolMsgs) && (numSolutions > receivedSolMsgs)) {
						for (;;) {
							MPI_Test(&receivedSol, &receivedSolMsg, &unusedStat);
							if (receivedSolMsg) {
								printf("%s\n", outSols);
								receivedSolMsgs++;
								break;
							}
						}
					}
					while (numSolutions > receivedSolMsgs) {
						/*printf("stuck waiting for solutions\n");*/
						MPI_Recv(outSols, solutionSize, MPI_CHAR, MPI_ANY_SOURCE, 4, MPI_COMM_WORLD, &unusedStat);
						printf("%s\n", outSols);
						receivedSolMsgs++;
					}
				}

			}
			else {
				for (rowCount = 0; rowCount < boardSize; rowCount++) {
					MPI_Irecv(&ready, 1, MPI_INT, MPI_ANY_SOURCE, 1, MPI_COMM_WORLD, &receivedReadyReq);
					/* Loop forever until I receive any ready request */
					for(;;) {
						MPI_Test(&receivedReadyReq, &receivedReadyMsg, &readyStat);
						if (receivedReadyMsg) {
							MPI_Isend(&rowCount, 1, MPI_INT, readyStat.MPI_SOURCE, 0, MPI_COMM_WORLD, &unusedReq);
							break;
						}
					}
				}
				
				for (rowCount = 0; rowCount < boardSize; rowCount++) {
					MPI_Recv(&returnedNumSolutions, 1, MPI_UNSIGNED, MPI_ANY_SOURCE, 3, MPI_COMM_WORLD, &unusedStat);
					numSolutions += returnedNumSolutions;
				}
			}

			/* Tell all the worker ranks to quit */
			for (rankNum = 1; rankNum < totalRanks; rankNum++) {
				/* printf("Telling rank %d to terminate\n", rankNum); */
				MPI_Isend(&ready, 1, MPI_INT, rankNum, 2, MPI_COMM_WORLD, &unusedReq);
			}
		}
		else {
			/* Receive terminate message */
			MPI_Irecv(&allDone, 1, MPI_INT, 0, 2, MPI_COMM_WORLD, &receivedTerminateReq);
			while (!allDone) {
				/* Tell rank 0 that I'm ready for work */
				MPI_Isend(&ready, 1, MPI_INT, 0, 1, MPI_COMM_WORLD, &unusedReq);
				/* Receive work messages */
				MPI_Irecv(&myWorkRow, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, &receivedWorkReq);
				/* Loop forever until I've got work or I receive a terminate message */
				for (;;) {
					MPI_Test(&receivedWorkReq, &receivedWorkMsg, &unusedStat);
					if (receivedWorkMsg) {
						numSolutions = 0;
						solve(myWorkRow);
						MPI_Isend(&numSolutions, 1, MPI_UNSIGNED, 0, 3, MPI_COMM_WORLD, &unusedReq);
						break;
					}
					MPI_Test(&receivedTerminateReq, &receivedTerminateMsg, &unusedStat);
					if (receivedTerminateMsg) {
						/* printf("%d : received terminate message, allDone = %d\n", myRank, allDone); */
						break;
					}
				}
			}
		}
	}
	else {
		/* Single process */
		for (myWorkRow = 0; myWorkRow < boardSize; myWorkRow++) {
			solve(myWorkRow);
		}
	}
	return;
}

int main(int argc, char *argv[]) {
	int opt, rc;

	rc = MPI_Init(&argc,&argv);
	if (rc != MPI_SUCCESS) {
    	printf ("Error starting MPI program. Terminating.\n");
     	MPI_Abort(MPI_COMM_WORLD, rc);
	}
	MPI_Comm_size(MPI_COMM_WORLD, &totalRanks);
	MPI_Comm_rank(MPI_COMM_WORLD, &myRank);

	/*  http://www.opengroup.org/onlinepubs/000095399/functions/getopt.html */
	while ((opt = getopt(argc, argv, "qn:")) != -1) {
		if (opt == 'q') {
			printSolutions = 0;
		}
		else if (opt == 'n') {
			boardSize = atoi(optarg);
			if ((boardSize <= 0) || (boardSize > 25)) {
				usage();
				MPI_Finalize();
				return -1;
			}
		}
		else {
			printf("Unrecognized option %c\n\n", optarg);
			usage();
			MPI_Finalize();
			return -1;
		}
	}
	/* printf ("Number of tasks = %d My rank = %d\n", totalRanks, myRank); */

	setup();

	if (myRank == 0) {
		printf("%d total solutions\n", numSolutions);
		if (solutions[boardSize-1] == numSolutions) {
			printf("success\n");
		}
		else {
			printf("failure\n");
		}
	}
	MPI_Finalize();
	return 0;
}
